# 版本: 1.0.0
# 作者: 陈振玺
# 功能: 统计FASTQ文件中reads序列的出现次数

# 所有包导入必须放在文件头部
import argparse
import gzip
from collections import Counter
from pathlib import Path
from Bio import SeqIO
from typing import Dict, Tuple


class FastqSequenceCounter:
    """
    FASTQ序列计数器
    统计FASTQ文件中reads序列的出现次数
    """
    
    def __init__(self, input_file: str):
        """
        初始化FastqSequenceCounter
        
        Args:
            input_file: 输入的FASTQ文件路径（支持压缩格式）
        """
        self.input_file = input_file
        self.sequence_counts = Counter()
    
    def is_gzipped(self, file_path: str) -> bool:
        """
        判断文件是否为gzip压缩格式
        
        Args:
            file_path: 文件路径
            
        Returns:
            bool: 是否为gzip压缩文件
        """
        with open(file_path, 'rb') as f:
            return f.read(2) == b'\x1f\x8b'
    
    def count_sequences(self) -> Dict[str, int]:
        """
        统计FASTQ文件中序列的出现次数
        
        Returns:
            Dict[str, int]: 序列及其出现次数的字典
            
        Raises:
            Exception: 处理文件时发生错误
        """
        try:
            # 判断文件是否为压缩格式
            if self.is_gzipped(self.input_file):
                # 处理压缩的FASTQ文件
                with gzip.open(self.input_file, 'rt') as handle:
                    for record in SeqIO.parse(handle, 'fastq'):
                        sequence = str(record.seq)
                        self.sequence_counts[sequence] += 1
            else:
                # 处理普通的FASTQ文件
                with open(self.input_file, 'r') as handle:
                    for record in SeqIO.parse(handle, 'fastq'):
                        sequence = str(record.seq)
                        self.sequence_counts[sequence] += 1
            
            print(f"成功处理文件: {self.input_file}")
            print(f"总reads数: {sum(self.sequence_counts.values())}")
            print(f"唯一序列数: {len(self.sequence_counts)}")
            
            return dict(self.sequence_counts)
            
        except FileNotFoundError:
            raise Exception(f"文件不存在: {self.input_file}")
        except Exception as e:
            raise Exception(f"处理FASTQ文件失败: {str(e)}")
    
    def get_top_sequences(self, top_n: int = 10) -> list:
        """
        获取出现次数最多的序列
        
        Args:
            top_n: 返回前N个最常见的序列
            
        Returns:
            list: 包含(序列, 计数)元组的列表
        """
        return self.sequence_counts.most_common(top_n)
    
    def save_results(self, output_file: str, top_n: int = None):
        """
        保存结果到CSV文件
        
        Args:
            output_file: 输出文件路径
            top_n: 只保存前N个最常见的序列，如果为None则保存所有
            
        Raises:
            Exception: 保存文件时发生错误
        """
        try:
            # 确定要保存的数据
            if top_n is not None:
                sequences_to_save = self.get_top_sequences(top_n)
                print(f"保存前 {top_n} 个最常见的序列")
            else:
                sequences_to_save = self.sequence_counts.most_common()
                print(f"保存所有 {len(self.sequence_counts)} 个唯一序列")
            
            # 写入CSV文件
            with open(output_file, 'w', encoding='utf-8') as f:
                # 写入CSV头部
                f.write("sequence,count,length\n")
                # 写入数据
                for sequence, count in sequences_to_save:
                    sequence_length = len(sequence)
                    f.write(f"\"{sequence}\",{count},{sequence_length}\n")
            
            print(f"结果已保存到: {output_file}")
            
        except Exception as e:
            raise Exception(f"保存结果失败: {str(e)}")
    
    def print_statistics(self):
        """
        打印统计信息
        """
        total_reads = sum(self.sequence_counts.values())
        unique_sequences = len(self.sequence_counts)
        
        print("\n" + "=" * 50)
        print("统计信息")
        print("=" * 50)
        print(f"总reads数: {total_reads:,}")
        print(f"唯一序列数: {unique_sequences:,}")
        print(f"重复率: {((total_reads - unique_sequences) / total_reads * 100):.2f}%")
        
        # 显示前10个最常见的序列
        top_10 = self.get_top_sequences(10)
        if top_10:
            print("\n前10个最常见的序列:")
            print("-" * 50)
            for i, (sequence, count) in enumerate(top_10, 1):
                percentage = (count / total_reads) * 100
                # 如果序列太长，只显示前50个字符
                display_seq = sequence if len(sequence) <= 50 else sequence[:47] + "..."
                print(f"{i:2d}. {display_seq}\t{count:,}\t({percentage:.2f}%)")
        
        print("=" * 50)


def main():
    """
    主函数
    """
    parser = argparse.ArgumentParser(
        description='统计FASTQ文件中reads序列的出现次数',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
    # 统计所有序列
    fastq-sequence-counter -i input.fastq -o output.txt
    
    # 只统计前100个最常见的序列
    fastq-sequence-counter -i input.fastq -o output.txt --top 100
    
    # 处理压缩的FASTQ文件
    fastq-sequence-counter -i input.fastq.gz -o output.txt
        """
    )
    
    parser.add_argument(
        '-i', '--input',
        required=True,
        help='输入的FASTQ文件路径（支持.gz压缩格式）'
    )
    
    parser.add_argument(
        '-o', '--output',
        required=True,
        help='输出文件路径，格式：序列<tab>计数'
    )
    
    parser.add_argument(
        '--top',
        type=int,
        help='只输出前N个最常见的序列（可选）'
    )
    
    parser.add_argument(
        '--no-stats',
        action='store_true',
        help='不打印统计信息'
    )
    
    args = parser.parse_args()
    
    try:
        # 创建计数器对象
        counter = FastqSequenceCounter(args.input)
        
        # 统计序列
        counter.count_sequences()
        
        # 打印统计信息
        if not args.no_stats:
            counter.print_statistics()
        
        # 保存结果
        counter.save_results(args.output, args.top)
        
    except Exception as e:
        print(f"错误: {e}")
        exit(1)


if __name__ == "__main__":
    main()